#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
@author: liang kang
@contact: gangkanli1219@gmail.com
@time: 1/2/18 11:51 AM
@desc: prepare data
"""
import argparse
import os

import cv2
import numpy as np

from crowdcounting.data.revolver import CrowdCountingRevolver
from crowdcounting.data.tools import check_all_crowd_counting
from crowdcounting.data.tools import check_dat_and_img_file
from crowdcounting.data.tools import decode_dat_file
from utils.io import read_text_file
from utils.io import write_text_file
from utils.list import create_file_list
from utils.list import create_train_val_list


def parse_args():
    parser = argparse.ArgumentParser(description='Prepare lists txt file for dataset')
    parser.add_argument('--output', dest='output', help='输出的根目录', default='', type=str)
    parser.add_argument('--root', dest='root', help='数据的根目录', default='', type=str)
    parser.add_argument('--postfix', dest='postfix', help='读取的文件和生成的文件前缀', default='', type=str)
    parser.add_argument('--prefix', dest='prefix', help='读取的文件和生成的文件后缀', default='', type=str)
    return parser.parse_args()


def write_img_and_points(data, output, paths):
    """
    将图像和点数据保存

    Parameters
    ----------
    data
    output
    paths

    Returns
    -------

    """
    # 用于记录对应图像没有人的子图
    with open(os.path.join(output, 'non.txt'), 'a') as file:
        print(paths, file=file)

    for _, angle in data.items():
        for _, seat in angle.items():
            for key, info in seat.items():
                cv2.imwrite(os.path.join(output, '{}.jpg'.format(key)), info['image'])
                points = info['point']

                # 如果子图无人，保存对应key
                if len(points) == 0:
                    with open(os.path.join(output, 'non.txt'), 'a') as file:
                        print(key, file=file)
                    continue

                # 将坐标点写入TXT文件
                points = list(map(lambda x: str(x[0]) + ' ' + str(x[1]), points))
                write_text_file(os.path.join(output, '{}.txt'.format(key)), points)


def convert_dat_to_txt(file_list, output, split='&!&'):
    """
    按固定位置和角度旋转切割图像，并保存子图和其对应的点

    Parameters
    ----------
    file_list
    output
    split

    Returns
    -------

    """
    # 定义位置
    params = {'crop': {'edge': {
        'seat': ((0, 480), (80, 480), (160, 480), (240, 480), (320, 480)),
        'size': (320, 320)}}}

    if not os.path.exists(output):
        os.makedirs(output)
    for idx, paths in enumerate(file_list):
        print(paths)
        dat_path, img_path = paths.split(split)
        out_root = os.path.join(output, '{:04}'.format(idx // 100),
                                '{:08}'.format(idx))
        if not os.path.exists(out_root):
            os.makedirs(out_root)

        # 解析dat文件生成坐标点列表
        points = decode_dat_file(dat_path)

        # TODO 交换坐标，令宽度在前
        points = np.asarray(points)[:, (1, 0)]

        # 执行分割旋转
        revolver = CrowdCountingRevolver(img_path.strip(), points, revolve_angle=6, params=params)
        revolver.update()

        write_img_and_points(revolver.output, out_root, paths)


def raw_dat_to_point(param):
    """

    生成train和val数据集
    转换原始数据为可用于创建tf-record文件的数据

    Parameters
    ----------
    param

    Returns
    -------

    """
    # 定义路径
    root = param.root
    output = param.output
    output_train = os.path.join(output, 'train')
    output_val = os.path.join(output, 'val')

    # 将全部原始数据分离为train 和val
    file_list, _ = create_file_list(root, filtering=check_dat_and_img_file, params='&!&')
    create_train_val_list(file_list, output, param.prefix + '_', '_raw_' + param.postfix)
    train_list = read_text_file(os.path.join(output, '{}_train_raw_{}.txt'.format(param.prefix, param.postfix)))
    val_list = read_text_file(os.path.join(output, '{}_val_raw_{}.txt'.format(param.prefix, param.postfix)))

    # 转换原始数据
    convert_dat_to_txt(train_list, output_train)
    convert_dat_to_txt(val_list, output_val)

    # 创建输出子图的train和val文件列表
    create_file_list(output_train, os.path.join(output, '{}_train_{}.txt'.format(param.prefix, param.postfix)),
                     filtering=check_all_crowd_counting, params='&!&')
    create_file_list(output_val, os.path.join(output, '{}_val_{}.txt'.format(param.prefix, param.postfix)),
                     filtering=check_all_crowd_counting, params='&!&')


def count_the_number(param):
    """
    统计所有子图中的最大和最小人数

    Parameters
    ----------
    param

    Returns
    -------

    """

    def _count_file_lines(lists, split, file_path=''):
        """
        通过txt文件行数计算当前列表中的文件表示的子图及其人数
        返回已经得到的最大值和最小值，对应的文件

        Parameters
        ----------
        lists
        split
        file_path

        Returns
        -------

        """
        min_line, max_line = 10000, 0
        for idx, files in enumerate(lists):
            if idx % 1000 == 0:
                print('count {} of {}'.format(idx, len(lists)))
            txt, _ = files.split(split)
            with open(txt, 'rb') as file:
                txt_list = len(file.readlines())
            min_line = min(min_line, txt_list)
            if max_line < txt_list:
                file_path = txt
            max_line = max(max_line, txt_list)
        return min_line, max_line, file_path

    output = param.output
    train_list = read_text_file(os.path.join(output, '{}_train_{}.txt'.format(param.prefix, param.postfix)))
    min_number, max_number, file_name = _count_file_lines(train_list, '&!&')
    val_list = read_text_file(os.path.join(output, '{}_val_{}.txt'.format(param.prefix, param.postfix)))
    new_min, new_max, file_name = _count_file_lines(val_list, '&!&', file_name)
    print(min(min_number, new_min))
    print(max(max_number, new_max))
    print(file_name)


def filter_people(param, number=24, split='&!&'):
    """
    过滤掉子图中人数少于 `number` 的部分

    Parameters
    ----------
    param
    number
    split

    Returns
    -------

    """

    def _filter_num(file_name):
        txt, _ = file_name.split(split)
        with open(txt, 'rb') as file:
            num = len(file.readlines())
        return num >= number

    output = param.output
    train_list = read_text_file(os.path.join(output, '{}_train_{}.txt'.format(param.prefix, param.postfix)))
    train_list = list(filter(_filter_num, train_list))
    write_text_file(os.path.join(output, '{}_train_{}_final.txt'.format(param.prefix, param.postfix)), train_list)
    val_list = read_text_file(os.path.join(output, '{}_val_{}.txt'.format(param.prefix, param.postfix)))
    val_list = list(filter(_filter_num, val_list))
    write_text_file(os.path.join(output, '{}_val_{}_final.txt'.format(param.prefix, param.postfix)), val_list)


if __name__ == '__main__':
    args = parse_args()
    # count_the_number(args)
    # raw_dat_to_point(args)
    filter_people(args)
